#include "rphoton.hpp"

template<int _Vaxis> int comparePhotonImpacts(const void *pa, const void *pb)
{
	const rpPhotonImpactData *a = reinterpret_cast<const rpPhotonImpactData *>(pa);
	const rpPhotonImpactData *b = reinterpret_cast<const rpPhotonImpactData *>(pb);

	if(a->location[_Vaxis] > b->location[_Vaxis]) return -1;
	if(a->location[_Vaxis] < b->location[_Vaxis]) return 1;
	return 0;
}

void InsertPhotonClusterIntoLeaf(rpPhotonCatalog *node, rpPhotonImpactData *photons, unsigned int offset, unsigned int numPhotons)
{
	node->data.numPhotons = numPhotons;
	node->data.offset = offset;
	node->data.isLeaf = true;
}

void InsertPhotonClusterIntoTree(rpPhotonCatalog *node, rpPhotonImpactData *photons, unsigned int offset, unsigned int numPhotons, int currentAxisSort)
{
	terVec3 mins, maxs;

	if(numPhotons <= config.photonsPerCatalogCluster)
	{
		InsertPhotonClusterIntoLeaf(node, photons, offset, numPhotons);
		return;
	}

	rpPhotonImpactData *localPhotons = photons + offset;
	mins = localPhotons[0].location;
	maxs = localPhotons[0].location;

	for(unsigned int i=1;i<numPhotons;i++)
	{
		terVec3 loc = localPhotons[i].location;
		for(int j=0;j<3;j++)
		{
			if(loc[j] < mins[j]) mins[j] = loc[j];
			if(loc[j] > maxs[j]) maxs[j] = loc[j];
		}
	}

	unsigned int axisChoices[3] = { 0, 1, 2 };
	int primaryAxis;

	terVec3 vol = maxs-mins;

	if(vol[0] > vol[1])
	{
		if(vol[0] > vol[2])
			primaryAxis = 0;
		else
			primaryAxis = 2;
	}
	else
	{
		if(vol[1] > vol[2])
			primaryAxis = 1;
		else
			primaryAxis = 2;
	}

	axisChoices[0] = primaryAxis;
	axisChoices[primaryAxis] = 0;

	for(int axisIndex=0;axisIndex<3;axisIndex++)
	{
		int axis = axisChoices[axisIndex];

		terVec3 dir(0.0f, 0.0f, 0.0f);
		dir[axis] = 1.0f;

		// These might be in a particular order from a previous sort
		// If they are, then the sort process doesn't need to be repeated
		if(currentAxisSort != axis)
		{
			if(axis == 0)
				qsort(localPhotons, numPhotons, sizeof(rpPhotonImpactData), comparePhotonImpacts<0>);
			else if(axis == 1)
				qsort(localPhotons, numPhotons, sizeof(rpPhotonImpactData), comparePhotonImpacts<1>);
			else
				qsort(localPhotons, numPhotons, sizeof(rpPhotonImpactData), comparePhotonImpacts<2>);
		}
		currentAxisSort = axis;

		float axisDist;
		
		if(numPhotons % 2 == 1)
			axisDist = localPhotons[numPhotons/2].location[axis];
		else
		{
			axisDist = (localPhotons[numPhotons/2].location[axis] + localPhotons[numPhotons/2-1].location[axis]) * 0.5f;
		}

		Math::Plane dividingPlane(axisDist, dir);

		// Find the first photon that is on the backside
		unsigned int divisionPoint;
		for(divisionPoint=numPhotons/2;divisionPoint<numPhotons;divisionPoint++)
			if(dividingPlane.PointSide(localPhotons[divisionPoint].location) == Math::Plane::Back_k)
				break;

		if(divisionPoint != numPhotons)
		{
			// This works
			node->data.planeNormal = dividingPlane.Normal();
			node->data.planeDist = dividingPlane.Distance();
			node->data.offset = offset;
			node->data.numPhotons = numPhotons;
			node->data.isLeaf = false;
	
			node->children[0] = new rpPhotonCatalog();
			node->children[1] = new rpPhotonCatalog();

			InsertPhotonClusterIntoTree(node->children[0], photons, offset, divisionPoint, currentAxisSort);
			InsertPhotonClusterIntoTree(node->children[1], photons, offset+divisionPoint, numPhotons-divisionPoint, currentAxisSort);
			return;
		}
	}

	printf("Failed to find a division plane\n");
	// No good axis could be found, tightly-packed or degenerate cluster of some sort, store it as a leaf
	InsertPhotonClusterIntoLeaf(node, photons, offset, numPhotons);
}

void WriteCatalogNode(terFileStream *catalogFile, const rpPhotonCatalog *node)
{
	catalogFile->Write(&node->data, sizeof(rpPhotonCatalogData));
	if(!node->data.isLeaf)
	{
		WriteCatalogNode(catalogFile, node->children[0]);
		WriteCatalogNode(catalogFile, node->children[1]);
	}
}

static terFileSystem *collateFS;
static terString collateBasePath;

static void CollatePhotonChart(int chartNum)
{
	terVector<rpPhotonImpactData> photons;
	int sizeOfImpact = sizeof(rpPhotonImpactData);

	printf("Collating chart %i...\n", chartNum);
	terSafePointer<rpPhotonCatalog> rootNode = new rpPhotonCatalog();

	terSafePointer<terFileStream> photonFile = collateFS->Open(collateBasePath + "_photons/chart" + terString(chartNum) + ".photons");

	int prevOutOffset;

	unsigned int numPhotons = photonFile->Size() / sizeof(rpPhotonImpactData);

	if(numPhotons == 0)
		return;	// Nothing to do

	if(numPhotons > config.maxPhotonsPerManifold)
	{
		photons.Alloc(config.maxPhotonsPerManifold);

		// Too many photons in this catalog
		float energyMultiplier = ((float)numPhotons) / ((float)config.maxPhotonsPerManifold);

		int previousIndex = -1;
		for(unsigned int i=0;i<numPhotons;i++)
		{
			__int64 newIndex64 = ((__int64)i) * ((__int64)config.maxPhotonsPerManifold) / ((__int64)numPhotons);
			int newIndex = (int)newIndex64;
			rpPhotonImpactData data;

			photonFile->Read(&data, sizeof(data));

			if(newIndex != previousIndex)
			{
				// Can include this
				photons[newIndex] = data;
			}
			previousIndex = newIndex;
		}

		numPhotons = config.maxPhotonsPerManifold;
	}
	else
	{
		photons.Alloc(numPhotons);
		photonFile->Read(photons.Recast<void>(), sizeof(rpPhotonImpactData) * numPhotons);
	}

	InsertPhotonClusterIntoTree(rootNode, photons, 0, numPhotons, -1);

	// Write sorted photons
	photonFile.Discard();
	photonFile = collateFS->OpenUnprotectedWrite(collateBasePath + "_photons/chart" + terString(chartNum) + ".photons");
	photonFile->Write(photons, sizeof(rpPhotonImpactData) * numPhotons);

	terSafePointer<terFileStream> catalogFile = collateFS->OpenUnprotectedWrite(collateBasePath + "_photons/chart" + terString(chartNum) + ".catalog");
	WriteCatalogNode(catalogFile, rootNode);
}

class rpCollateWorkUnit : public Threading::WorkUnit
{
	terSafePointer<terMutex> _mtx;
	class Threading::Job *_parent;
	int _chartID;

public:
	void SetChartID(int chartID)
	{
		_chartID = chartID;
	}

	rpCollateWorkUnit()
	{
		_mtx = ThreadManager->CreateMutex();
		_chartID = 0;
	}

	void Work()
	{
		CollatePhotonChart(_chartID);

		_mtx->Release();
	}

	bool TryAcquire()
	{
		return _mtx->TryAcquire();
	}
};


class rpCollateChartJob : public Threading::DecomposableJob
{
public:
	rpCollateChartJob(Threading::JobManager *parent, unsigned int numCharts)
		: DecomposableJob(parent, numCharts)
	{
	}

	~rpCollateChartJob()
	{
	}

	Threading::WorkUnit *NewWorkUnit()
	{
		return new rpCollateWorkUnit();
	}

	void PrepareWorkUnit(Threading::WorkUnit *bru)
	{
		rpCollateWorkUnit *wu = (rpCollateWorkUnit *)bru;
		wu->SetChartID(_currentUnit);
	}
};


void CollatePhotonFiles(terFileSystem *fs, const terString &basePath, Threading::JobManager *jm)
{
	int objID = 0;
	collateFS = fs;
	collateBasePath = basePath;

	terSafePointer<Threading::Job> job = new rpCollateChartJob(jm, geometry._numCharts);
	jm->AddJob(job);

	while(jm->ActiveJobs())
		ThreadManager->SleepThread(1000);
}
